import os
import os.path as osp
import platform
import argparse
from tqdm import tqdm
import numpy as np
import tflite_runtime.interpreter as tflite
import torch
from utils import get_test_dataset, set_input, get_output


EDGETPU_SHARED_LIB = {
  'Linux': 'libedgetpu.so.1',
  'Darwin': 'libedgetpu.1.dylib',
  'Windows': 'edgetpu.dll'
}[platform.system()]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default=osp.join('data', 'qat_edgetpu.tflite'))
    parser.add_argument('--dataset', type=str, choices=['mnist', 'cifar10'], default='mnist')
    args = parser.parse_args()

    interpreter = tflite.Interpreter(
        model_path=args.model,
        experimental_delegates=[
            tflite.load_delegate(EDGETPU_SHARED_LIB)
        ])
    interpreter.allocate_tensors()

    _, test_loader = get_test_dataset(args.dataset)

    for i, (image, label) in enumerate(test_loader, 1):
        image = (image.permute(0, 2, 3, 1)[0].numpy() * 255).astype(np.uint8)
        set_input(
            interpreter=interpreter,
            image=image)
        interpreter.invoke()
        y = get_output(interpreter=interpreter)
        print('label={}, y={}'.format(label, np.argmax(y)))


if __name__ == '__main__':
    main()
